#pragma once

// https://blog.csdn.net/fulianzhou/article/details/125186185

#include <cstdint>
#include <string>
#include <fstream>
#include <vector>
#include <memory>
#include <iostream>
#include <arpa/inet.h>
#include <map>
#include <array>
#include <cstring>
// pcap 文件头

#pragma pack(1)
struct PcapFileHeader
{
    uint32_t magic; // 0xa1b2c3d4
    uint16_t version_major;
    uint16_t version_minor;
    int32_t thiszone;  // gmt to local correction
    uint32_t sigfigs;  // accuracy of timestamps
    uint32_t snaplen;  // max length saved portion of each pkt
    uint32_t linktype; // data link type (LINKTYPE_*)
};

struct PcapTimeVal
{
    uint32_t tv_sec;  // seconds
    uint32_t tv_usec; // microseconds
};

// pcap 数据包头
struct PcapPacketHeader
{
    PcapTimeVal ts;  // time stamp
    uint32_t caplen; // length of portion present
    uint32_t len;    // length this packet (off wire)
};

// ethernet 数据帧头
struct EthernetFrameHeader
{
    uint8_t dst_mac[6];
    uint8_t src_mac[6];
    uint16_t type;
};

// ipv4 数据包头
struct Ipv4PacketHeader
{
#if __BYTE_ORDER == __LITTLE_ENDIAN
    uint8_t header_length : 4;
    uint8_t version : 4;
#elif __BYTE_ORDER == __BIG_ENDIAN
    uint8_t version : 4;
    uint8_t header_length : 4;
#endif
    uint8_t tos;
    uint16_t total_length;
    uint16_t id;
    union
    {
        struct
        {
#if __BYTE_ORDER == __LITTLE_ENDIAN
            uint16_t fragment_offset : 13;
            uint16_t flags : 3;
#elif __BYTE_ORDER == __BIG_ENDIAN
            uint16_t flags : 3;
            uint16_t fragment_offset : 13;
#endif
        };
        uint16_t all;
    } flags_fragment_offset;

    uint8_t ttl;
    uint8_t protocol;
    uint16_t checksum;
    uint32_t src_ip;
    uint32_t dst_ip;

    void ntoh()
    {
        total_length = ntohs(total_length);
        id = ntohs(id);
        flags_fragment_offset.all = ntohs(flags_fragment_offset.all);
        checksum = ntohs(checksum);
        src_ip = ntohl(src_ip);
        dst_ip = ntohl(dst_ip);
    }
    void hton()
    {
        total_length = htons(total_length);
        id = htons(id);
        flags_fragment_offset.all = htons(flags_fragment_offset.all);
        checksum = htons(checksum);
        src_ip = htonl(src_ip);
        dst_ip = htonl(dst_ip);
    }
};

// udp 数据包头
struct UdpPacketHeader
{
    uint16_t src_port;
    uint16_t dst_port;
    uint16_t length;
    uint16_t checksum;
    void hton()
    {
        src_port = htons(src_port);
        dst_port = htons(dst_port);
        length = htons(length);
        checksum = htons(checksum);
    }
    void ntoh()
    {
        src_port = ntohs(src_port);
        dst_port = ntohs(dst_port);
        length = ntohs(length);
        checksum = ntohs(checksum);
    }
};

// tcp 数据包头
struct TcpPacketHeader
{
    uint16_t src_port;
    uint16_t dst_port;
    uint32_t seq_num;
    uint32_t ack_num;
    uint8_t reserved : 4;
    uint8_t header_length : 4;
    uint8_t flags;
    uint16_t window_size;
    uint16_t checksum;
    uint16_t urgent_pointer;
    void ntoh()
    {
        src_port = ntohs(src_port);
        dst_port = ntohs(dst_port);
        seq_num = ntohl(seq_num);
        ack_num = ntohl(ack_num);
        window_size = ntohs(window_size);
        checksum = ntohs(checksum);
        urgent_pointer = ntohs(urgent_pointer);
    }
};

#pragma pack()

class PcapPaser
{
    // 只保留 dest == 239.16.18.16 的数据

    // 源服务器有：172.19.255.19:32949, 172.19.255.18:56349
    // 目标地址：239.16.18.16:9011/9012/9013

private:
    const std::string m_file_path;
    std::ifstream m_file;
    std::vector<uint8_t> m_content;

    std::map<uint16_t, std::vector<std::vector<uint8_t>>> m_ip_fragment;

public:
    PcapPaser(const std::string &file_path) : m_file_path(file_path)
    {
    }
    ~PcapPaser()
    {
        if (m_file.is_open())
            m_file.close();
    }

    bool open()
    {
        m_file.open(m_file_path, std::ios::binary);
        if (!m_file.is_open())
            return false;
        std::unique_ptr<PcapFileHeader> file_header(new PcapFileHeader);
        m_file.read((char *)file_header.get(), sizeof(PcapFileHeader));
        if (file_header->magic != 0xa1b2c3d4)
        {
            return false;
        }
        std::cout << "magic: " << std::hex << file_header->magic << "\t"
                  << "version: " << std::dec << file_header->version_major
                  << "." << file_header->version_minor
                  << std::endl;

        return true;
    }

    // void readAll()
    // {
    //     // 获取文件大小
    //     m_file.seekg(0, std::ios::end);
    //     std::streampos end_pos = m_file.tellg();
    //     m_file.seekg(0, std::ios::beg);
    //     std::streampos beg_pos = m_file.tellg();

    //     size_t file_size = end_pos - beg_pos;
    //     m_content.resize(file_size - sizeof(PcapFileHeader));
    //     // 将除了文件头的内容读入内存
    //     m_file.seekg(sizeof(PcapFileHeader), std::ios::beg);
    //     m_file.read((char *)m_content.data(), m_content.size());
    // }

    // void parse()
    // {
    // }

    bool isEOF() const
    {
        return m_file.eof();
    }

    std::vector<uint8_t> getIP()
    {
        if (isEOF())
        {
            return {};
        }

        std::array<uint8_t, sizeof(PcapPacketHeader) + sizeof(EthernetFrameHeader) + sizeof(Ipv4PacketHeader)> header;
        m_file.read((char *)header.data(), header.size());

        PcapPacketHeader *pcap_header = (PcapPacketHeader *)header.data();

        EthernetFrameHeader *ethernet_header = (EthernetFrameHeader *)(header.data() + sizeof(PcapPacketHeader));
        if (ntohs(ethernet_header->type) != 0x0800) // 不是IP报文
        {
            return {};
        }

        Ipv4PacketHeader *ipv4_header = (Ipv4PacketHeader *)(header.data() + sizeof(PcapPacketHeader) + sizeof(EthernetFrameHeader));

        // 判断源IP是否为 172.19.205.41
        if (ntohl(ipv4_header->src_ip) != 0xac13cd29)
        {
            m_file.seekg(ntohs(ipv4_header->total_length) - sizeof(Ipv4PacketHeader), std::ios::cur);
            return {};
        }
        // 判断目标ip是否为 172.19.205.202
        if (ntohl(ipv4_header->dst_ip) != 0xac13cdca) // 不是目标ip
        {
            m_file.seekg(ntohs(ipv4_header->total_length) - sizeof(Ipv4PacketHeader), std::ios::cur);
            return {};
        }

        // 不是TCP报文
        if (ipv4_header->protocol != 0x06)
        {
            m_file.seekg(ntohs(ipv4_header->total_length) - sizeof(Ipv4PacketHeader), std::ios::cur);
            return {};
        }

        std::vector<uint8_t> res;
        res.resize(ntohs(ipv4_header->total_length));
        memcpy(res.data(), header.data() + sizeof(PcapPacketHeader) + sizeof(EthernetFrameHeader), sizeof(Ipv4PacketHeader));
        m_file.read((char *)res.data() + sizeof(Ipv4PacketHeader), res.size() - sizeof(Ipv4PacketHeader));
        return std::move(res);
    }

    std::vector<uint8_t> getUDP()
    {
        if (isEOF())
        {
            return {};
        }

        std::array<uint8_t, sizeof(PcapPacketHeader) + sizeof(EthernetFrameHeader) + sizeof(Ipv4PacketHeader) + sizeof(UdpPacketHeader)> header;
        m_file.read((char *)header.data(), header.size());

        PcapPacketHeader *pcap_header = (PcapPacketHeader *)header.data();

        EthernetFrameHeader *ethernet_header = (EthernetFrameHeader *)(header.data() + sizeof(PcapPacketHeader));
        if (ntohs(ethernet_header->type) != 0x0800) // 不是IP报文
        {
            return {};
        }

        Ipv4PacketHeader *ipv4_header = (Ipv4PacketHeader *)(header.data() + sizeof(PcapPacketHeader) + sizeof(EthernetFrameHeader));
        ipv4_header->ntoh();

        // 判断目标ip是否为 239.16.18.16
        if (ipv4_header->dst_ip != 0xef101210) // 不是目标ip
        {
            return {};
        }

        if (ipv4_header->protocol != 0x11) // 不是UDP报文
        {
            return {};
        }

        std::vector<uint8_t> udp_data; // 存储udp报文体

        // 如果是分片报文
        if (ipv4_header->flags_fragment_offset.flags == 0x1) // 分片
        {
            // m_ip_fragment[ntohs(ipv4_header->id)].emplace_back();
            return {};
        }
        else if (ipv4_header->flags_fragment_offset.flags == 0x0) // 分片的最后一片
        {
            return {};
        }
        else
        {
            udp_data.resize(
                ipv4_header->total_length - sizeof(Ipv4PacketHeader) -
                sizeof(UdpPacketHeader));

            m_file.read((char *)udp_data.data(), udp_data.size());
        }

        return std::move(udp_data);
    }
};
